import math
from gettext import translation

import pandas as pd
import torch
from torch import nn
from torch.nn.functional import dropout

import dltools
from NLP_selfAtttention import PositionalEncoding, num_steps, num_hiddens, num_heads


# 基于位置的前置网络
# position wise feed-forward net
class PositionWiseFFN(nn.Module):
    def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs, **kwargs):
        super().__init__(**kwargs)
        self.dense1 = nn.Linear(ffn_num_input, ffn_num_hiddens)
        self.relu = nn.ReLU()
        self.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)

    def forward(self, X):
        return self.dense2(self.relu(self.dense1(X)))

ffn = PositionWiseFFN(6, 4, 8)
ffn.eval()
X = torch.ones((2, 3, 6)) #(batchSize, num_steps, vocab_size)
# pytorch Linear默认输入数据是2维度，输入更多维度时会把前面的维度合并，保留最后一个维度，当做2维进行计算，输出时会进行还原
res = ffn(X)
print(f"res.shape:{res.shape:}")


class AddNorm(nn.Module):
    def __init__(self, normalized_shape, dropout, **kwargs):
        super().__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)
        # bn是batch之间的数据进行归一化，ln是batch内的数据做归一化
        self.ln = nn.LayerNorm(normalized_shape)

    def forward(self, X, Y):
        # X Y 形状要相同，简单的按位置相加
        return self.ln(self.dropout(Y) + X)


add_norm = AddNorm([6, 8], 0.2)
add_norm.eval()
X = torch.ones((2, 6, 8), dtype = torch.float32)
Y = torch.ones((2, 6, 8), dtype = torch.float32)
res = add_norm(X, Y)
print(res)


# 编码器block
class EncoderBlock(nn.Module):
    def __init__(self, key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, ffn_num_outputs, num_heads,  dropout=0.0, use_bias=False, **kwargs):
        super().__init__(**kwargs)
        self.attention = dltools.MultiHeadAttention(key_size, query_size, value_size, num_hiddens, num_heads, dropout, use_bias)
        self.add_norm1 = AddNorm(norm_shape, dropout)
        self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens, ffn_num_outputs)
        self.add_norm2 = AddNorm(norm_shape, dropout)

    def forward(self, X, valid_lens):
        temp = self.add_norm1(X, self.attention(X, X, X, valid_lens))
        return self.add_norm2(temp, self.ffn(temp))


X = torch.ones((2, 100, 24)) #(batchSize, num_steps, vocab_size)
valid_lens = torch.tensor([3, 2])
encoder_blk = EncoderBlock(24, 24, 24,24, [100, 24], 24, 48, 24,8, 0.1)
encoder_blk.eval()
res = encoder_blk(X, valid_lens)
print(f"res.shape:{res.shape}")


# 编码器
class TransformerEncoder(dltools.Encoder):
    def __init__(self, vocab_size, key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, num_heads, num_layers, dropout, use_bias= False, **kwargs):
        super().__init__(**kwargs)
        self.num_hiddens = num_hiddens
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        self.pos_encoding = PositionalEncoding(num_hiddens, dropout)
        self.blks = nn.Sequential()
        for i in range(num_layers):
            self.blks.add_module("block" + str(i), EncoderBlock(key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens, num_hiddens, num_heads,  dropout, use_bias))

    def forward(self, X, valid_lens, *args):
        # 先对embedding数据进行缩放，有助于数据收敛。
        print("-------")
        print(f"X.shape:{X.shape}")
        print(f"self.embedding(X).shape:{self.embedding(X).shape}")
        print("-------")
        X= self.embedding(X) * math.sqrt(self.num_hiddens)
        # pos_encoding操作已经包含了 X+P的操作
        X = self.pos_encoding(X)
        self._attention_weights = [None] * len(self.blks)

        for i, blk in enumerate(self.blks):
            X = blk(X, valid_lens)
            self._attention_weights[i] = blk.attention.attention.attention_weights
        return X

    @property
    def attention_weights(self):
        return self._attention_weights

encoder_1 = TransformerEncoder(200, 24, 24, 24, 24, [100, 24], 24, 48,  8, 2, 0.1)
encoder_1.eval()
X = torch.ones((2, 100), dtype=torch.long)
valid_lens = torch.tensor([3, 2])
res = encoder_1(X, valid_lens)
print("------")
print(f"res.shape:{res.shape}")

# decoder block
class DecoderBlock(nn.Module):
    def __init__(self, key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,  num_heads,   dropout=0.0, use_bias=False,  i=0, **kwargs):
        super().__init__(**kwargs)
        self.i = i
        self.attention1 = dltools.MultiHeadAttention(key_size, query_size, value_size, num_hiddens, num_heads, dropout, use_bias)
        self.addNorm1 = AddNorm(norm_shape, dropout)
        self.attention2 = dltools.MultiHeadAttention(key_size, query_size, value_size, num_hiddens, num_heads, dropout, use_bias)
        self.addNorm2 = AddNorm(norm_shape, dropout)
        self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens, num_hiddens)
        self.addNorm3 = AddNorm(norm_shape, dropout)

    def forward(self, X, state):
        enc_output, enc_valid_lens = state[0], state[1]
        # print(f"state[2]:{state[2]}")
        if state[2][self.i] is None:
            # 训练
            key_values = X
        else:
            # 预测：需要把前面时刻预测得到的信息和当前block的输出信息拼接到一起
            key_values = torch.cat((state[2][self.i], X), dim=1)
        state[2][self.i] = key_values
        if self.training:
            dev_valid_lens = torch.arange(1, X.shape[1] + 1, device=X.device).repeat(X.shape[0], 1)
            print(f"dev_valid_lens:{dev_valid_lens}")
        else:
            dev_valid_lens = None

        # 自注意力
        X2 = self.attention1(X, key_values, key_values, dev_valid_lens)
        Y = self.addNorm1(X, X2)
        Y2 = self.attention2(Y, enc_output, enc_output, enc_valid_lens)
        Z = self.addNorm2(Y, Y2)
        return self.addNorm3(Z, self.ffn(Z)), state

print("----------")
decoder_blk = DecoderBlock(24, 24, 24, 24, [100, 24], 24, 48, 8, 0.5, i = 0)
decoder_blk.eval()
X = torch.ones((2, 100, 24))
valid_lens = torch.tensor([3, 2])
state = [encoder_blk(X, valid_lens), valid_lens, [None]]
res1, res2 = decoder_blk(X, state)
print(f"res1:{res1}")
print(f"res2[2]:{res2[2]}")

class TransformerDecoder(dltools.Decoder):
    def __init__(self, vocab_size,  key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,  num_heads, num_layers, dropout=0.0, use_bias=False, **kwargs):
        super().__init__(**kwargs)
        self.num_hiddens = num_hiddens
        self.num_layers = num_layers
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        self.pos_encoding = PositionalEncoding(num_hiddens, dropout)
        self.blks = nn.Sequential()
        for i in range(num_layers):
            self.blks.add_module("block"+str(i), DecoderBlock(key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,  num_heads,   dropout, use_bias,  i, **kwargs))

        self.dense = nn.Linear(num_hiddens, vocab_size)

    def init_state(self, enc_outputs, enc_valid_lens, *args):
        return [enc_outputs, enc_valid_lens, [None] * self.num_layers]

    def forward(self, X, state):
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        self._attention_weights = [[None] * len(self.blks) for _ in range(2)]

        for i, blk in enumerate(self.blks):
            X, state = blk(X, state)
            self._attention_weights[0][i] = blk.attention1.attention.attention_weights
            self._attention_weights[1][i] = blk.attention2.attention.attention_weights

        return self.dense(X), state

    @property
    def attention_weights(self):
        return self._attention_weights


num_hiddens, num_layers, dropout, batch_size, num_steps = 32, 2, 0.1, 64, 10
lr, num_epochs, device = 0.005, 200, dltools.try_gpu()
ffn_num_inputs, ffn_num_hiddens, num_heads = 32, 64, 4
keySize, query_size, value_size = 32, 32, 32
norm_shape = [32]

train_iter, src_vocab, tgt_vocab = dltools.load_data_nmt(batch_size, num_steps)

encoder_2 = TransformerEncoder(len(src_vocab), keySize, query_size, value_size, num_hiddens, norm_shape, ffn_num_inputs, ffn_num_hiddens, num_heads, num_layers, dropout)
decoder_2 = TransformerDecoder(len(tgt_vocab), keySize, query_size, value_size, num_hiddens, norm_shape, ffn_num_inputs, ffn_num_hiddens, num_heads, num_layers, dropout)
net = dltools.EncoderDecoder(encoder_2, decoder_2)
dltools.train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)

# 开始预测
engs = ["go .", "i lost .", "he's calm .", "i'm home ."]
fras = ["va !", "j'ai perdu .", "il est calme .", "je suis chez moi ."]

for eng, fra in zip(engs, fras):
    translation =  dltools.predict_seq2seq(net, eng, src_vocab, tgt_vocab, num_steps, device)
    print(f"{eng} => {translation}, bleu:{dltools.bleu(translation[0], fra, k=2):.3f}")

